"""
yolo-ui-2.py
create by cuiyu 2025-08-07
说明：
 * YOLOv11（Ultralytics）可视化界面GUI设计，基于pyside6，单文件即插即用
 * https://blog.csdn.net/qq_67105081/article/details/144974796
 * 代码来源：上面地址，可以直接运行。

用法:
  yolo-ui-2.py

更新:
2025-08-07:
"""


import cv2,sys
import numpy as np
from PySide6.QtGui import QIcon, QImage, QPixmap
from PySide6.QtWidgets import QApplication, QMainWindow, QLabel, QVBoxLayout, QWidget, QPushButton, QHBoxLayout, QMessageBox, QFileDialog, QGroupBox
from PySide6.QtCore import Qt,QEvent,QTimer
from ultralytics import YOLO
 
class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
 
        self.xiaolian_ui()
        self.model = None
        self.timer = QTimer()
        self.cap = None
        self.result = None
 
    def xiaolian_ui(self):
        self.setFixedSize(1200, 630)
        self.setWindowTitle('@author：笑脸惹桃花')
        self.setWindowIcon(QIcon("icon.png"))
        self.move(400,300)
        centralWidget = QWidget(self)
        self.setCentralWidget(centralWidget)
        mainlayout = QVBoxLayout(centralWidget)
        layout = QHBoxLayout()
        self.label1 = QLabel(self)
        self.label1.setMinimumSize(580, 550)
        self.label1.setStyleSheet('border:3px solid #6950a1; background-color: black;')
        layout.addWidget(self.label1)
        self.label2 = QLabel(self)
        self.label2.setMinimumSize(580, 550)
        self.label2.setStyleSheet('border:3px solid #6950a1; background-color: black;')
        layout.addWidget(self.label2)
        mainlayout.addLayout(layout)
        # 界面下半部分： 输出框 和 按钮
        groupbox = QGroupBox(self)
        bottomlayout = QVBoxLayout(groupbox)
 
        mainlayout.addWidget(groupbox)
 
        hbox_buttons = QHBoxLayout()
        self.load_model_btn = QPushButton('👆选择模型')
        self.load_model_btn.setFixedSize(120, 30)
        self.load_model_btn.clicked.connect(self.load_model)
        hbox_buttons.addWidget(self.load_model_btn)
        self.img_detect_btn = QPushButton('🖼️选择图片')
        self.img_detect_btn.setFixedSize(120, 30)
        self.img_detect_btn.setEnabled(False)
        self.img_detect_btn.clicked.connect(self.select_image)
        hbox_buttons.addWidget(self.img_detect_btn)
        self.display_btn = QPushButton("🔍显示检测物体")
        self.display_btn.clicked.connect(self.show_detected_objects)
        self.display_btn.setEnabled(False)
        self.display_btn.setFixedSize(120, 30)
        hbox_buttons.addWidget(self.display_btn)
        self.stop_detect_btn = QPushButton('🛑停止')
        self.stop_detect_btn.setFixedSize(120, 30)
        self.stop_detect_btn.setEnabled(False)
        self.stop_detect_btn.clicked.connect(self.stop_detect)
        hbox_buttons.addWidget(self.stop_detect_btn)
        self.exit_btn = QPushButton('❌退出')
        self.exit_btn.setFixedSize(120, 30)
        self.exit_btn.clicked.connect(self.close)
        hbox_buttons.addWidget(self.exit_btn)
        bottomlayout.addLayout(hbox_buttons)
 
    def load_model(self):
        model_path, _ = QFileDialog.getOpenFileName(self, "选择模型文件", filter='*.pt')
        if model_path:
            self.model = YOLO(model_path)
        else:
            print("请重选模型！")
        self.img_detect_btn.setEnabled(True)
        self.stop_detect_btn.setEnabled(True)
 
    def select_image(self):
        if self.timer.isActive():
            self.timer.stop()
        if self.cap is not None:
            self.cap.release()
        image_path, fileType = QFileDialog.getOpenFileName(self, "选择图片文件", filter='*.jpg *.png *.bmp')
        if image_path:
            img = cv2.imread(image_path)
            self.detect_image(img)
 
    def show_detected_objects(self):
        frame = self.result
        if frame:
            det_info = []
            class_ids = frame[0].boxes.cls
            class_names_dict = frame[0].names
            for class_id in class_ids:
                class_name = class_names_dict[int(class_id)]
                det_info.append(class_name)
            if det_info:
                object_count = len(det_info)
                object_info = f"识别到的物体总个数：{object_count}\n"
                object_dict = {}
                for obj in det_info:
                    if obj in object_dict:
                        object_dict[obj] += 1
                    else:
                        object_dict[obj] = 1
                sorted_objects = sorted(object_dict.items(), key=lambda x: x[1], reverse=True)
 
                for obj_name, obj_count in sorted_objects:
                    object_info += f"{obj_name}: {obj_count}\n"
                self.show_message_box("识别结果", object_info)
            else:
                self.show_message_box("识别结果", "未检测到物体")
 
    def stop_detect(self):
        if self.timer.isActive():
            self.timer.stop()
        if self.cap is not None:
            self.cap.release()
            self.cap = None
        self.result = None
        img = cv2.cvtColor(np.zeros((580, 550), np.uint8), cv2.COLOR_BGR2RGB)
        img = QImage(img.data, img.shape[1], img.shape[0], QImage.Format_RGB888)
        self.label1.setPixmap(QPixmap.fromImage(img))
        self.label2.setPixmap(QPixmap.fromImage(img))
        self.display_btn.setEnabled(False)
 
    def close(self):
        if self.cap is not None:
            self.cap.release()
            self.cap = None
        if self.timer.isActive():
            self.timer.stop()
        exit()
 
    def detect_image(self, img):
        if self.model is not None:
            frame = img
            # Use model for inference
            results = self.model.predict(frame)
            image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            height1, width1, channel1 = frame.shape
            bytesPerLine1 = 3 * width1
            qimage1 = QImage(image_rgb.data, width1, height1, bytesPerLine1, QImage.Format_RGB888)
            pixmap1 = QPixmap.fromImage(qimage1)
            self.label1.setPixmap(pixmap1.scaled(self.label1.size(), Qt.AspectRatioMode.IgnoreAspectRatio))
            annotated_image = results[0].plot()
            annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)  # 转换为 RGB
            height2, width2, channel2 = annotated_image.shape
            bytesPerLine2 = 3 * width2
            qimage2 = QImage(annotated_image.data, width2, height2, bytesPerLine2, QImage.Format_RGB888)
            pixmap2 = QPixmap.fromImage(qimage2)
            self.label2.setPixmap(pixmap2.scaled(self.label2.size(), Qt.AspectRatioMode.IgnoreAspectRatio))
            self.result = results
            self.display_btn.setEnabled(True)
 
    def show_message_box(self, title, message):
        msg_box = QMessageBox(self)
        msg_box.setWindowTitle(title)
        msg_box.setText(message)
        msg_box.exec()
 
    def closeEvent(self, event: QEvent):
        # 关闭应用时清除缓存
        self.cap.release()
        self.timer.stop()
        event.accept()
 
if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())